from tkinter import *
from tkinter import ttk
from tkinter import messagebox
import torch

Attr = dict(有无手术经历={'有': 1, '无': 2, '未知': 0}, 年龄={'成年': 1, '未成年': 2, '未知': 0}, 直肠温度={}, 脉搏={}, 呼吸频率={},
            四肢温度={'正常': 1, '温暖': 2, '冷': 3, '冰冷': 4, '未知': 0}, 周血管脉搏={'正常': 1, '过大': 2, '过小': 3, '未知': 0},
            黏膜色泽={'正常粉': 1, '亮粉': 2, '淡粉': 3, '淡紫': 4, '亮红': 5, '暗紫': 6, '未知': 0},
            血管充填时间={'小于3秒': 1, '大于3秒': 2, '未知': 0}, 马痛苦等级={'无痛苦': 1, '消沉': 2, '间歇微疼': 3, '间歇剧疼': 4, '持续疼痛': 5, '未知': 0},
            内脏蠕动情况={'蠕动有力': 1, '正常': 2, '蠕动衰退': 3, '无蠕动': 4, '未知': 0},
            腹胀情况={'无': 1, '轻微': 2, '适度': 3, '一直': 4, '未知': 0}, 鼻腔气体={'无': 1, '轻微': 2, '重': 3, '未知': 0},
            鼻胃反流={'无': 1, '大于1L': 2, '小于1L': 3, '未知': 0}, 反流液PH={}, 粪便情况={'正常': 1, '增加': 2, '衰减': 3, '无': 4, '未知': 0},
            腹腔情况={'正常': 1, '其他': 2, '肠胃积粪': 3, '小肠扩张': 4, '大肠扩张': 5, '未知': 0}, 红细胞占比={}, 蛋白质占比={},
            腹腔穿刺外观={'清晰': 1, '模糊': 2, '血色': 3, '未知': 0}, 腹腔穿刺总蛋白={})

# 添加下拉框
top = Tk(className="马疝病预测")
for i, lable in enumerate(Attr):
    globals()[f'label{i}'] = Label(top, text=lable, background="WhiteSmoke", font=20)
    globals()[f'label{i}'].grid(column=(i % 3), row=(i // 3)*2, padx=5, pady=5)
    globals()[f'cmb{i}'] = ttk.Combobox(top, font=20)
    globals()[f'cmb{i}']["value"] = tuple(Attr[lable].keys())
    if Attr[lable] != {}:
        globals()[f'cmb{i}'].current(0)
    globals()[f'cmb{i}'].grid(column=(i % 3), row=(i // 3)*2 + 1, padx=5, pady=5)

# Pre的钩子函数
def PreClick():
    def Getdata():
        data = []
        for i, key in enumerate(Attr):
            cmb_str = globals()[f'cmb{i}'].get()
            if cmb_str == "":
                data.append(0)
            else:
                if Attr[key] != {}:
                    data.append(Attr[key][cmb_str])
                else:
                    data.append(float(cmb_str))
        return data

    model = torch.load("./checkpoint/model.pkl").eval()
    data_in = torch.tensor(Getdata()).unsqueeze(0).float()
    out = model(data_in)
    if out >= 0.5:
        messagebox.showinfo("诊断结果", "根据您输入的病理学特征分析\r\n您的马匹状态为:\r\n*******************************\r\n                "
                                    "     健康\r\n*******************************\r\n健康分数：" + str(float(out*100))[:5])
    else:
        messagebox.showinfo("诊断结果", "根据您输入的病理学特征分析\r\n您的马匹状态为:\r\n*******************************\r\n                 "
                                    "    疾病\r\n*******************************\r\n健康分数：" + str(float(out*100))[:5])


def window_init():
    top.geometry("600x550")
    top.resizable(False, False)
    #top.iconbitmap("./myicon.ico")
    # 预测按键
    button_pre = Button(text="预测", height=2, width=20, command=PreClick)
    button_pre.grid(column=2, pady=20)
    return top


if __name__ == "__main__":
    top = window_init()
    top.mainloop()
